/**
 * 
 */
package edu.umd.clip.lm.model.decoding;

import java.util.*;

import edu.umd.clip.lm.model.decoding.TrellisSlice.SliceItem;
import edu.umd.clip.lm.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class LinkQueue extends AbstractQueue<SliceItem> implements Cloneable {
	private SortedRingBuffer<SliceItem> elements;
	private int maxElements;
	private double probThreshold;
	private double totalProb; 
	private double cumulativeProbThreshold;
	private double consumedProbability;
	private int consumedElements;
	private double prunedProbability;
	private HashSet<SliceItem> prunedElements;
	
	private static final int INIT_BUFFER_SIZE = 8;
	
	protected static final Comparator<SliceItem> comparator = new Comparator<SliceItem>() {
		@Override
		public int compare(SliceItem arg0, SliceItem arg1) {
			// reverse order
			return - (int) Math.signum(arg0.getProb() - arg1.getProb());
		}
	};

	protected LinkQueue(Collection<SliceItem> items, int maxElements, double threshold, double cumulativeProbFraction) {
		SliceItem[] arr = items.toArray(new SliceItem[0]);
		Arrays.sort(arr, comparator);
		this.maxElements = maxElements;
		totalProb = 0;
		prunedProbability = 0;
		double[] partSum = new double[arr.length];
		for(int i=0; i<arr.length; ++i) {
			SliceItem item = arr[i];
			totalProb += item.getProb();
			partSum[i] = totalProb;
		}
		prunedElements = new HashSet<SliceItem>();
		
		cumulativeProbThreshold = cumulativeProbFraction * totalProb;
		if (arr.length > maxElements) {
			int cutOffPos = Arrays.binarySearch(partSum, cumulativeProbThreshold);
			if (cutOffPos < 0) {
				cutOffPos = - cutOffPos - 1;
			}
			cutOffPos += 10; // make it a little over the threshold
			if (cutOffPos > maxElements) {
				cutOffPos = maxElements;
			}
			if (cutOffPos < arr.length * 9 / 10 || arr.length > maxElements) {
				for(int i=cutOffPos; i<arr.length; ++i) {
					addPrunedElement(arr[i]);					
				}
				arr = Arrays.copyOf(arr, cutOffPos);
				prunedProbability = totalProb - partSum[cutOffPos-1];
			}
		}
		double bestProb = partSum[0];
		probThreshold = threshold * bestProb;
		
		elements = new SortedRingBuffer<SliceItem>(arr, Math.max(arr.length, INIT_BUFFER_SIZE), comparator);
		
		consumedProbability = 0;
		consumedElements = 0;
		if (Decoder.DEBUG > 1 && prunedProbability > 0.0) {
			System.out.printf("LinkQueue: initial pruned %g, totalProb=%g, elements=%d\n", 
					prunedProbability, totalProb, elements.size());
		}

	}

	private LinkQueue() {
	}
	
	public void consumeProbability(double prob) {
		consumedProbability += prob;
		++consumedElements;
	}
	
	public boolean isDone() {
		return consumedElements >= maxElements || consumedProbability >= cumulativeProbThreshold;
	}
	
	/* (non-Javadoc)
	 * @see java.util.AbstractCollection#size()
	 */
	@Override
	public int size() {
		return elements.size();
	}
	/* (non-Javadoc)
	 * @see java.util.AbstractCollection#iterator()
	 */
	@Override
	public Iterator<SliceItem> iterator() {
		return elements.iterator();
	}
	/* (non-Javadoc)
	 * @see java.util.Queue#offer(java.lang.Object)
	 */
	@Override
	public boolean offer(SliceItem e) {
		if (e == null) {
			throw new IllegalArgumentException();
		}
		if (e.getProb() > probThreshold) {
			if (elements.isFull()) {
				if (elements.size() < maxElements) {
					int newSize = Math.min(elements.size()*2, maxElements);
					elements.setSize(newSize);
					return elements.offer(e);
				}
				SliceItem last = elements.peekLast();
				if (comparator.compare(e, last) < 0) {
					// remove the lowest probability element
					elements.pollLast();
					addPrunedElement(last);
					prunedProbability += last.getProb();
					
					if (Decoder.DEBUG > 1) {
						System.out.printf("LinkQueue: adding %g, pruning %g, totalProb=%g, totalPruned = %g\n", 
								e.getProb(), last.getProb(), totalProb, prunedProbability);
					}
				} else {
					// reject the item
					//prunedProbability += e.getProb();
					return false;
				}
			}
			return elements.offer(e);
		}
		return true;
	}
	/* (non-Javadoc)
	 * @see java.util.Queue#peek()
	 */
	@Override
	public SliceItem peek() {
		return elements.peek();
	}
	/* (non-Javadoc)
	 * @see java.util.Queue#poll()
	 */
	@Override
	public SliceItem poll() {
		SliceItem item = elements.poll();
		//System.out.println("queue.poll() = " + item);
		return item;
	}

	public double getTotalProb() {
		return totalProb;
	}

	protected void addPrunedElement(SliceItem item) {
		//System.out.println("pruned "+item.toString());
		prunedElements.add(item);
	}
	
	protected void addPrunedProbability(double prob) {
		prunedProbability += prob;
	}
	
	public double getPrunedProbability() {
		return prunedProbability;
	}

	public HashSet<SliceItem> getPrunedElements() {
		return prunedElements;
	}
	
	public void clearPrunedElements() {
		prunedElements.clear();
	}

	@SuppressWarnings("unchecked")
	@Override
	public Object clone() {
		LinkQueue copy = new LinkQueue();
		copy.consumedElements = consumedElements;
		copy.consumedProbability = consumedProbability;
		copy.cumulativeProbThreshold = cumulativeProbThreshold;
		copy.elements = (SortedRingBuffer<SliceItem>)elements.clone();
		copy.maxElements = maxElements;
		copy.probThreshold = probThreshold;
		if (prunedElements != null) {
			copy.prunedElements = new HashSet<SliceItem>(prunedElements);
		}
		copy.prunedProbability = prunedProbability;
		copy.totalProb = totalProb;
		return copy;
	}

	@Override
	public boolean isEmpty() {
		return elements.isEmpty();
	}
	
}
